# -*- coding: utf-8 -*-
import argparse
import os
from pathlib import Path
from pdf2image import convert_from_path
import logging

logging.basicConfig(
    format="[{asctime}]-{levelname}-{message}",
    style="{",
    datefmt="%Y-%m-%d %H:%M",
    level=logging.INFO # Add a default logging level
)


def pdf_to_png(pdf_path, output_dir, dpi=300, first_n_pages=None):
    """将单个 PDF 文件转换为 PNG 图片。

    Args:
        pdf_path (str): PDF 文件的路径。
        output_dir (str): 保存 PNG 图片的目录。
        dpi (int, optional): 输出 PNG 图片的 DPI。默认为 300。
        first_n_pages (int, optional): 如果指定，则只转换前 N 页。默认为 None (转换所有页)。
    """
    try:
        # Ensure output_dir exists
        Path(output_dir).mkdir(parents=True, exist_ok=True)

        images = convert_from_path(pdf_path, dpi=dpi)
        pdf_name = Path(pdf_path).stem # Get the name without suffix for base filename
        # output_path_base = Path(output_dir) / Path(pdf_path).stem # This part was slightly off for handling directories

        # Construct the base output path correctly within the output_dir
        output_path_base = Path(output_dir) / pdf_name

        for i, image in enumerate(images):
            if first_n_pages is not None and i >= first_n_pages:
                break
            output_file = f"{output_path_base}_page_{i + 1}.png"
            image.save(output_file, "PNG")
        logging.info(f"Successfully converted '{pdf_path}'.\n") # Added output_dir to log
    except Exception as e:
        logging.error(f"转换 '{pdf_path}' 失败: {e}")

def process_path(input_path, output_dir=None, dpi=300, n_pages = None): # Changed default to None
    """处理给定的输入路径，如果是文件则转换，如果是文件夹则递归处理所有 PDF 文件。

    Args:
        input_path (str): 输入的文件或文件夹路径。
        output_dir (str, optional): 保存 PNG 图片的根目录。如果为 None，则输出到输入文件或文件夹的同级目录。
        dpi (int, optional): 输出 PNG 图片的 DPI。默认为 300。
        n_pages (int, optional): 对于单个 PDF 文件，指定转换的前 N 页。默认为 None (转换所有页)。
    """
    input_p = Path(input_path)

    # Determine the actual output directory based on the default or user input
    if output_dir is None:
        if input_p.is_file():
            # If input is a file, output to its parent directory
            output_p = input_p.parent
        elif input_p.is_dir():
            # If input is a directory, output to the input directory itself (subfolders will be mirrored)
            output_p = input_p
        else:
             logging.warn(f"参数`input_path {input_path}` 不是一个 PDF 文件或文件夹。")
             return # Exit if the input is invalid
    else:
        # If output_dir is provided, use it
        output_p = Path(output_dir)

    # output_p.mkdir(parents=True, exist_ok=True) # mkdir is now handled inside pdf_to_png for single files

    if input_p.is_file() and input_p.suffix.lower() == ".pdf":
        # For a single file, output_p is its parent directory
        pdf_to_png(str(input_p), str(output_p), dpi=dpi, first_n_pages=n_pages)
    elif input_p.is_dir():
        # For a directory, iterate and mirror the structure
        for root, _, files in os.walk(input_path):
            relative_root = Path(root).relative_to(input_path)
            # Current output directory is relative to the determined output_p
            current_output_dir = output_p / relative_root
            # current_output_dir.mkdir(parents=True, exist_ok=True) # mkdir is now handled inside pdf_to_png

            for file in files:
                if file.lower().endswith(".pdf"):
                    pdf_file_path = os.path.join(root, file)
                    # The output directory for each PDF is within the mirrored structure
                    pdf_to_png(pdf_file_path, str(current_output_dir), dpi=dpi, first_n_pages=n_pages) # Added n_pages here
    else:
         logging.warn(f"参数`input_path {input_path}` 不是一个 PDF 文件或文件夹。")


def main():
    parser = argparse.ArgumentParser(description="将 PDF 文件转换为 PNG 图片。")
    parser.add_argument("input_path", help="要处理的 PDF 文件或文件夹路径。")
    # Changed default to None, and updated help text
    parser.add_argument("-o", "--output_dir", default=None,
                         help="保存 PNG 图片的根目录。如果未指定，则输出到输入文件或文件夹的同级目录。")
    parser.add_argument("-d", "--dpi", type=int, default=300,
                         help="输出 PNG 图片的 DPI (默认为 300)。", required=False)
    # Removed default=3 for pages, making it None by default unless specified
    parser.add_argument("-p", "--pages", type=int, default=None,
                         help="单个pdf文件转化时，指定转化页数 (默认为 None，转化所有页)。", required=False)
    args = parser.parse_args()

    process_path(input_path=args.input_path,
                 output_dir=args.output_dir,
                 dpi=args.dpi,
                 n_pages=args.pages) # Pass n_pages to process_path

if __name__ == "__main__":
    main()